package com.yx.vectordb.impl.lucene;

import lombok.extern.slf4j.Slf4j;
import org.apache.lucene.analysis.standard.StandardAnalyzer;
import org.apache.lucene.document.*;
import org.apache.lucene.index.*;
import org.apache.lucene.search.*;
import org.apache.lucene.store.*;

import java.io.IOException;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Slf4j
public class LuceneVectorEngine {
    private final Directory directory;
    private final IndexWriter writer;
    private final int dimensions;
    private final StandardAnalyzer analyzer;

    public LuceneVectorEngine(String indexPath, int dimensions) throws IOException {
        this.dimensions = dimensions;
        this.analyzer = new StandardAnalyzer();
        
        directory = FSDirectory.open(Paths.get(indexPath));
        IndexWriterConfig config = new IndexWriterConfig(analyzer);
        config.setMergePolicy(new LogByteSizeMergePolicy());
        writer = new IndexWriter(directory, config);
    }

    public void addVector(String id, float[] vector) throws IOException {
        if (vector.length != dimensions) {
            throw new IllegalArgumentException("Vector dimensions mismatch");
        }

        // 先删除同ID的旧数据
        deleteVector(id);

        Document doc = new Document();
        // 添加ID
        doc.add(new StringField("id", id, Field.Store.YES));
        // 添加向量
        doc.add(new KnnFloatVectorField("vector", vector));
        
        // 添加元数据
        writer.addDocument(doc);
        writer.commit();
    }

    // 修改 createSearchResults 方法以排除 _exact 后缀的字段
    private List<LuceneSearchResult> createSearchResults(IndexReader reader, TopDocs topDocs) throws IOException {
        List<LuceneSearchResult> results = new ArrayList<>();
        for (ScoreDoc scoreDoc : topDocs.scoreDocs) {
            Document doc = reader.storedFields().document(scoreDoc.doc);
            Map<String, String> metadata = new HashMap<>();
            
            for (IndexableField field : doc.getFields()) {
                String name = field.name();
                if (!name.equals("id") && !name.equals("vector") && !name.endsWith("_exact")) {
                    metadata.put(name, field.stringValue());
                }
            }
            
            results.add(new LuceneSearchResult(
                doc.get("id"),
                scoreDoc.score,
                metadata
            ));
        }
        return results;
    }


    public List<LuceneSearchResult> search(float[] queryVector, int k) throws IOException {
        if (queryVector.length != dimensions) {
            throw new IllegalArgumentException("Query vector dimensions mismatch");
        }

        writer.commit();
        
        try (IndexReader reader = DirectoryReader.open(directory)) {
            IndexSearcher searcher = new IndexSearcher(reader);
            Query knnQuery = new KnnFloatVectorQuery("vector", queryVector, k);
            {
                TopDocs topDocs = searcher.search(knnQuery, k);
                return createSearchResults(reader, topDocs);
            }
        }
    }

    

    public void deleteVector(String id) throws IOException {
        writer.deleteDocuments(new Term("id", id));
        writer.commit();
    }

    public void commit() throws IOException {
        writer.commit();
    }

    public void close() throws IOException {
        writer.close();
        directory.close();
    }

}